"""
Configuration classes for hook managers and related components.

This module provides configuration for module discovery, hook factories,
and the overall hook manager system.
"""

from dataclasses import dataclass, field, asdict
from typing import Optional, List, Dict, Any, Callable, Type, Union
import torch.nn as nn

from continuallearning.models.pefts.common.config.config import BasePEFTConfig


@dataclass
class ModuleDiscoveryConfig:
    """Configuration for module discovery strategies.

    Default Behaviour:
    - Discovers Linear layers by default (find_linear_layers=True)
    - Searches recursively through the entire model (recursive=True)
    - Excludes frozen modules (include_frozen=False)
    - No depth limit for searching (max_depth=None)
    - Common transformer patterns are pre-configured in __post_init__
    """

    _config_name = "module_discovery"

    # Type-based discovery
    target_types: Optional[List[Type[nn.Module]]] = None

    # Pattern-based discovery
    target_patterns: Optional[List[str]] = None
    pattern_type: str = "auto"  # "auto", "glob", "regex", "exact"
    stop_on_match: bool = True

    # Predicate-based discovery
    target_predicate: Optional[Callable[..., bool]] = None
    include_path_in_predicate: bool = False

    # Specific layer type discovery
    find_transformer_layers: bool = False
    find_linear_layers: bool = False
    find_conv_layers: bool = False

    # Transformer layer patterns
    transformer_name_patterns: Optional[List[str]] = None
    transformer_attribute_patterns: Optional[List[str]] = None

    # Additional filtering
    filter_predicate: Optional[Callable[..., bool]] = None

    # Discovery options
    cache_results: bool = True

    # 提供合理的默认值
    def __post_init__(self):
        """Set default patterns for transformer discovery."""
        if self.transformer_name_patterns is None:
            self.transformer_name_patterns = [
                "encoder.layer.*",
                "blocks.*",
                "encoders.*",
                "encoder_layers.*",
                "layers.*",
                "transformer.h.*",
                "decoder.layers.*",
            ]

        if self.transformer_attribute_patterns is None:
            self.transformer_attribute_patterns = [
                "attention",
                "attn",
                "self_attn",
                "self_attention",
                "mlp",
                "feed_forward",
                "ffn",
                "mha",
            ]

    def validate(self) -> None:
        """Validate module discovery configuration."""
        # Ensure at least one discovery method is specified
        discovery_methods = [
            self.target_types,
            self.target_patterns,
            self.target_predicate,
            self.find_transformer_layers,
            self.find_linear_layers,
            self.find_conv_layers,
        ]

        if not any(discovery_methods):
            raise ValueError(
                "At least one module discovery method must be specified. "
                "Options: target_types, target_patterns, target_predicate, "
                "find_transformer_layers, find_linear_layers, find_conv_layers"
            )

        valid_pattern_types = ["auto", "glob", "regex", "exact"]
        if self.pattern_type not in valid_pattern_types:
            raise ValueError(
                f"Invalid pattern_type: {self.pattern_type}. "
                f"Valid options: {valid_pattern_types}"
            )

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)


@dataclass
class RouterConfig:
    """Configuration for task routing."""

    _config_name = "router"

    # Router type and strategy
    num_experts: int = 1
    embed_dim: int = 768
    routing_strategy: str = "soft"
    weights_initializer: str = "kaiming_uniform_"

    # Routing parameters
    top_k: int = 2
    temperature: float = 1.0
    noise_std: float = 0.0

    # Model parameters (for learned routers)
    load_balancing_loss_weight: float = -1.0
    track_routing_history: bool = False

    def validate(self) -> None:
        """Validate router configuration."""
        if self.temperature <= 0:
            raise ValueError(f"temperature must be positive, got {self.temperature}")
        if self.noise_std < 0:
            raise ValueError(f"noise_std must be non-negative, got {self.noise_std}")

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)


@dataclass
class CombinerConfig:
    """Configuration for output combination.
    task_inter_weightedsum: no config
    concatenation:
        topk: int
    """

    _config_name = "combiner"

    # Combiner type
    combiner_type: str = "task_inter_weightedsum"

    # concatenation parameters
    topk: int = 10

    def validate(self) -> None:
        pass

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)


@dataclass
class HookFactoryConfig:
    """Configuration for hook factory."""

    _config_name = "hook_factory"

    # Factory identification
    factory_type: str = "base"

    # PEFT configuration
    hook_name: str = "base"
    hook_config: Optional[Union[BasePEFTConfig, Dict[str, Any]]] = None

    # Task-aware components (optional)
    router_name: Optional[str] = None
    router_config: Optional[Union[RouterConfig, Dict[str, Any]]] = None
    combiner_name: Optional[str] = None
    combiner_config: Optional[Union[CombinerConfig, Dict[str, Any]]] = None

    # Factory options
    lazy_init: bool = False
    share_weights: bool = False

    def __post_init__(self):
        """Process nested configurations."""
        # Convert dict to RouterConfig if needed
        if isinstance(self.router_config, dict):
            self.router_config = RouterConfig(**self.router_config)

        # Convert dict to CombinerConfig if needed
        if isinstance(self.combiner_config, dict):
            self.combiner_config = CombinerConfig(**self.combiner_config)

    def validate(self) -> None:
        """Validate hook factory configuration."""
        # If task-aware components are specified, ensure both are present
        if self.router_name and not self.combiner_name:
            raise ValueError("combiner_name must be specified when router_name is set")
        if self.combiner_name and not self.router_name:
            raise ValueError("router_name must be specified when combiner_name is set")

        # Validate nested configurations
        if isinstance(self.router_config, RouterConfig):
            self.router_config.validate()
        if isinstance(self.combiner_config, CombinerConfig):
            self.combiner_config.validate()

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)


@dataclass
class HookManagerConfig:
    """Complete configuration for hook manager."""

    _config_name = "hook_manager"

    # Module discovery configuration
    discovery: ModuleDiscoveryConfig = field(default_factory=ModuleDiscoveryConfig)

    # Hook factory configuration
    hook_factory: Union[str, HookFactoryConfig, Dict[str, Any]] = "default"

    # Manager behavior
    auto_attach: bool = True
    strict_mode: bool = False

    # Logging and debugging
    verbose: bool = False
    debug: bool = False
    log_level: str = "INFO"

    # Performance options
    lazy_init: bool = False
    batch_operations: bool = True
    parallel_init: bool = False

    # Resource management
    max_hooks: Optional[int] = None
    memory_efficient: bool = False

    def __post_init__(self):
        """Process nested configurations."""
        # Convert dict to ModuleDiscoveryConfig if needed
        if isinstance(self.discovery, dict):
            self.discovery = ModuleDiscoveryConfig(**self.discovery)
        # Convert dict/str to HookFactoryConfig if needed
        if isinstance(self.hook_factory, dict):
            self.hook_factory = HookFactoryConfig(**self.hook_factory)
        elif isinstance(self.hook_factory, str):
            # Simple string means default factory with that hook type
            self.hook_factory = HookFactoryConfig(
                factory_type="base", hook_name=self.hook_factory
            )

    def validate(self) -> None:
        """Validate hook manager configuration."""
        # Validate nested configurations
        self.discovery.validate()

        if isinstance(self.hook_factory, HookFactoryConfig):
            self.hook_factory.validate()

        # Validate manager-specific options
        if self.max_hooks is not None and self.max_hooks <= 0:
            raise ValueError(f"max_hooks must be positive, got {self.max_hooks}")

        valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
        if self.log_level.upper() not in valid_log_levels:
            raise ValueError(
                f"Invalid log_level: {self.log_level}. "
                f"Valid options: {', '.join(valid_log_levels)}"
            )

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)
